NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator#3798
NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator#3798singhharsh1708 wants to merge 2 commits into
Conversation
| ) | ||
| prob = if use_w_reuse | ||
| nlf_jac = let W = W | ||
| (z, p) -> convert(AbstractMatrix, W) |
There was a problem hiding this comment.
We shouldn't unconditionally convert.
|
good catch — split the closure so StaticWOperator returns its .W field directly, only WOperator path converts now |
350a7e8 to
0f20284
Compare
| end | ||
| else | ||
| let Ww = W | ||
| (z, p) -> convert(AbstractMatrix, Ww) |
There was a problem hiding this comment.
What about non-static opertators
|
for non-static |
d85b000 to
98bbc94
Compare
| end | ||
| else | ||
| let Ww = W | ||
| (z, p) -> Ww._concrete_form |
|
fixed — now hands the operator through directly |
833ee15 to
e414df1
Compare
|
Hi @ChrisRackauckas — quick nudge on this stack when you have a moment. All three PRs (#3795, #3796, #3798) have addressed the earlier review comments and are rebased on the latest master:
Once the next SciMLOperators release tags, I'll open a follow-up here to route OOP + SciMLOperator mass matrix through No rush — just wanted to make sure these weren't waiting on me. Thanks! |
|
Ran this branch against an OOP nonlinear problem and hit two defects in the W-reuse path — posting reproductions so they can be addressed: 1. The using OrdinaryDiffEqSDIRK, OrdinaryDiffEqNonlinearSolve, NonlinearSolve, ADTypes
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
vdp(u, p, t) = [u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1])]
prob = ODEProblem(vdp, [2.0, 0.0], (0.0, 1.0), [1.0e3])
integ = init(prob, TRBDF2(nlsolve = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff())), concrete_jac = true); reltol = 1e-8, abstol = 1e-10)
step!(integ)
# ERROR: MethodError: Cannot `convert` an object of type WOperator{false, ...} to an object of type Matrix{Float64}The 2. Even with the convert fixed, the Jacobian values are never refreshed — J is frozen at its build-time state for the entire integration. The update in The IIP counterpart (#3753) avoids this because Also worth folding in while touching this: pass |
e414df1 to
936de56
Compare
|
Thanks for the review @ChrisRackauckas-Claude — both defects reproduced and fixed. New commit Defect 1 (crash on convert): the OOP closure now returns the concrete matrix directly — Defect 2 (frozen J): added Ref indirection: the closure captures Tests: added two regressions at the end of
Locally: both new tests pass; existing |
bd5af11 to
d89e8d7
Compare
| nlf_jac = let W_ref = W_ref | ||
| (z, p) -> convert(AbstractMatrix, W_ref[]) | ||
| end | ||
| jac_prototype = copy(convert(AbstractMatrix, W)) |
There was a problem hiding this comment.
no it's still concretizing.
c752622 to
9cfc718
Compare
| W_ref = use_w_reuse ? Ref{typeof(W)}(W) : nothing | ||
| prob = if use_w_reuse | ||
| if W isa StaticWOperator | ||
| nlf_jac = let W_ref = W_ref | ||
| (z, p) -> W_ref[].W | ||
| end | ||
| NonlinearProblem( | ||
| NonlinearFunction{false, SciMLBase.FullSpecialize}( | ||
| nlf; jac = nlf_jac, jac_prototype = W.W | ||
| ), | ||
| copy(ztmp), nlp_params | ||
| ) | ||
| else | ||
| nlf_jac = let W_ref = W_ref | ||
| (z, p) -> convert(AbstractMatrix, W_ref[]) |
There was a problem hiding this comment.
OOP is purposefully not mutable. We want to make things fold away when static. This will violate that. I think it's simply not possible to do similar W-reuse in the OOP path without violating the principle of making this path non-mutating.
There was a problem hiding this comment.
For StaticWOperator the payload is an immutable SMatrix — reassigning nlcache.W[] re-points the Ref cell, the SMatrix itself is never mutated, so fold-away on the value side still holds. The mutation you're objecting to only shows up on the WOperator branch (copyto!(W.J, …)).
And picking NonlinearSolveAlg already forces a Matrix{Float64} jac cache inside construct_jacobian_cache — the OOP path isn't fold-away in that mode regardless of what this PR does. If Static-only reuse is still off the table given that, I'll close.
There was a problem hiding this comment.
I guess it's fine if there's a simplediffeq path that fully non-allocates with simplenonlinearsolve. But the conversion stuff is never okay.
That should get fixed upstream. Making everything dense is not a valid answer. |
9cfc718 to
f4cf1d3
Compare
|
Dropped WOperator branch. W-reuse only fires on StaticWOperator + AbstractSimpleNonlinearSolveAlgorithm — SMatrix rebuilt, Ref re-pointed, no convert. |
|
Why no WOperator? |
Drop the WOperator branch (concretizing/mutation-based).
SimpleNonlinearSolve algs land in NonlinearSolveNoInitCache which has no stats field.
b1b82a7 to
141a00b
Compare
|
@ChrisRackauckas Only StaticWOperator gives a real win in OOP: its For Can add a WOperator branch if you want the Vector-path coverage. |
The OOP
NonlinearSolveAlgbranch ofbuild_nlsolverbuildsJ, Wviabuild_J_Wand then discards both — theNonlinearProblemis built with nojac, andNonlinearSolveCachestoresnothingfor J/W. Every Newton iter inside the OOP path recomputes the Jacobian via AD/FD with no cross-step reuse.This is the OOP mirror of #3753 (which fixed the same gap on the IIP path for
WOperatorwith a concrete underlying J). The architecture is identical: detect a reusable W at build time, wire a closure jac into the innerNonlinearProblem, store J/W on the cache socompute_step!'s existingrecompute_jacobianflow picks them up.Scope
Enables W-reuse when
W isa Union{WOperator, StaticWOperator}(the case produced for OOP problems withconcrete_jac=true, a matrix-free linsolve likeKrylovJL_GMRES, or a SciMLOperatorjac_prototype). Other W types (plainMatrix,Factorization, genericAbstractSciMLOperator) fall through to the existing no-reuse default — so this is purely additive.Changes
utils.jlOOPbuild_nlsolverforNonlinearSolveAlg: detectuse_w_reuse, build a closure(z, p) -> convert(AbstractMatrix, W), pass it asjactoNonlinearFunction, storeJ, WonNonlinearSolveCache.newton.jlOOPinitialize!forNonlinearSolveAlg: whencache.W !== nothing, run the same first-call /new_W_dt_cutoff/ divergence logic the IIP path uses, callingupdate_coefficients!(cache.W, uprev, p, tstep; gamma = dtgamma)to refresh the WOperator in place and togglingcache.new_Wforcompute_step!.Local validation
Robertson, OOP,
ImplicitEuler(linsolve = KrylovJL_GMRES(), concrete_jac = true, ...):njacs = 0on the NS side reflects that the inner solver no longer computes J itself — the WOperator'supdate_coefficients!does it once per step instead.nw = 641over 1302 steps means W is reused across ~2 steps on average, matching the NLNewton modified-Newton cadence.